package main

import (
	"bufio"
	"context"
	"encoding/base64"
	"encoding/json"
	"io"
	"log"
	"net/http"
	"os"
	"sync"
	"time"

	"github.com/gorilla/websocket"
	"github.com/pion/rtpio"
	"github.com/pion/rtpio/pkg/rtpio"
	"github.com/pion/webrtc/v3"
)

type Signal struct {
	SDP       string `json:"sdp"`
	Type      string `json:"type"`
	Candidate string `json:"candidate,omitempty"`
}

var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool { return true },
}

func init() {
	log.SetFlags(log.LstdFlags | log.Lshortfile)
	log.Println("服务器启动...")
}

func main() {
	http.HandleFunc("/ws", handleWebSocket)
	log.Println("在 :8080 启动 HTTP 服务器")
	if err := http.ListenAndServe(":8080", nil); err != nil {
		log.Fatalf("HTTP 服务器启动失败: %v", err)
	}
}

func handleWebSocket(w http.ResponseWriter, r *http.Request) {
	log.Println("收到来自", r.RemoteAddr, "的 WebSocket 连接请求")
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("WebSocket 升级失败: %v", err)
		return
	}
	defer func() {
		log.Println("关闭 WebSocket 连接")
		conn.Close()
	}()

	// WebRTC 配置
	log.Println("使用 STUN 初始化 WebRTC 配置")
	config := webrtc.Configuration{
		ICEServers: []webrtc.ICEServer{
			{URLs: []string{"stun:stun.l.google.com:19302"}},
		},
	}

	// 创建媒体引擎并注册 H.264 和 Opus 编解码器
	log.Println("创建媒体引擎并注册编解码器")
	mediaEngine := &webrtc.MediaEngine{}
	if err := mediaEngine.RegisterCodec(webrtc.RTPCodecParameters{
		RTPCodecCapability: webrtc.RTPCodecCapability{
			MimeType:    webrtc.MimeTypeH264,
			ClockRate:   90000,
			SDPFmtpLine: "level-asymmetry-allowed=1;packetization-mode=1;profile-level-id=42e01f",
		},
		PayloadType: 96,
	}, webrtc.RTPCodecTypeVideo); err != nil {
		log.Fatalf("注册 H.264 编解码器失败: %v", err)
	}
	if err := mediaEngine.RegisterCodec(webrtc.RTPCodecParameters{
		RTPCodecCapability: webrtc.RTPCodecCapability{
			MimeType:    webrtc.MimeTypeOpus,
			ClockRate:   48000,
			Channels:    2,
			SDPFmtpLine: "minptime=10;useinbandfec=1",
		},
		PayloadType: 111,
	}, webrtc.RTPCodecTypeAudio); err != nil {
		log.Fatalf("注册 Opus 编解码器失败: %v", err)
	}

	// 创建 API 和 PeerConnection
	log.Println("创建 WebRTC API 和 PeerConnection")
	api := webrtc.NewAPI(webrtc.WithMediaEngine(mediaEngine))
	peerConnection, err := api.NewPeerConnection(config)
	if err != nil {
		log.Fatalf("创建 PeerConnection 失败: %v", err)
	}
	defer func() {
		log.Println("关闭 PeerConnection")
		peerConnection.Close()
	}()

	// 创建视频轨道
	log.Println("创建视频轨道")
	videoTrack, err := webrtc.NewTrackLocalStaticSample(webrtc.RTPCodecCapability{
		MimeType:  webrtc.MimeTypeH264,
		ClockRate: 90000,
	}, "video", "pion-video")
	if err != nil {
		log.Fatalf("创建视频轨道失败: %v", err)
	}

	// 创建音频轨道
	log.Println("创建音频轨道")
	audioTrack, err := webrtc.NewTrackLocalStaticSample(webrtc.RTPCodecCapability{
		MimeType:  webrtc.MimeTypeOpus,
		ClockRate: 48000,
		Channels:  2,
	}, "audio", "pion-audio")
	if err != nil {
		log.Fatalf("创建音频轨道失败: %v", err)
	}

	// 将轨道添加到 PeerConnection
	log.Println("将视频轨道添加到 PeerConnection")
	if _, err := peerConnection.AddTrack(videoTrack); err != nil {
		log.Fatalf("添加视频轨道失败: %v", err)
	}
	log.Println("将音频轨道添加到 PeerConnection")
	if _, err := peerConnection.AddTrack(audioTrack); err != nil {
		log.Fatalf("添加音频轨道失败: %v", err)
	}

	// 处理 RTCP 反馈
	go func() {
		log.Println("启动 RTCP 反馈循环")
		rtcpBuf := make([]byte, 1500)
		for {
			n, _, err := peerConnection.GetTransceivers()[0].Receiver().Read(rtcpBuf)
			if err != nil {
				log.Printf("RTCP 读取错误: %v", err)
				return
			}
			log.Printf("收到 %d 字节 RTCP 数据", n)
		}
	}()

	// 缓存 ICE 候选者
	pendingCandidates := make([]webrtc.ICECandidateInit, 0)
	var remoteDescSet bool

	// 处理 ICE 候选者
	peerConnection.OnICECandidate(func(c *webrtc.ICECandidate) {
		if c != nil {
			log.Printf("新的 ICE 候选者: %s", c.ToJSON().Candidate)
			candidate, _ := json.Marshal(c.ToJSON())
			if err := conn.WriteJSON(Signal{
				Candidate: string(candidate),
				Type:      "candidate",
			}); err != nil {
				log.Printf("发送 ICE 候选者失败: %v", err)
			}
		}
	})

	// 处理 ICE 连接状态变化
	peerConnection.OnICEConnectionStateChange(func(state webrtc.ICEConnectionState) {
		log.Printf("ICE 连接状态改变: %s", state.String())
		if state == webrtc.ICEConnectionStateFailed {
			log.Println("ICE 连接失败，关闭 PeerConnection")
			peerConnection.Close()
			conn.WriteJSON(Signal{Type: "error", SDP: "ICE 连接失败"})
		}
	})

	// 处理 WebSocket 消息
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	var wg sync.WaitGroup
	wg.Add(1)

	go func() {
		defer wg.Done()
		log.Println("启动 WebSocket 消息循环")
		for {
			_, msg, err := conn.ReadMessage()
			if err != nil {
				log.Printf("WebSocket 读取错误: %v", err)
				cancel()
				return
			}
			log.Printf("收到 WebSocket 消息: %s", string(msg))

			var signal Signal
			if err := json.Unmarshal(msg, &signal); err != nil {
				log.Printf("JSON 解析错误: %v", err)
				continue
			}

			switch signal.Type {
			case "offer":
				log.Println("收到客户端的提议")
				sdp, err := base64.StdEncoding.DecodeString(signal.SDP)
				if err != nil {
					log.Printf("SDP 解码错误: %v", err)
					continue
				}
				sdpStr := string(sdp)
				logSdp := sdpStr
				if len(sdpStr) > 200 {
					logSdp = sdpStr[:200] + "..."
				}
				log.Printf("解码的提议 SDP: %s", logSdp)

				if err := peerConnection.SetRemoteDescription(webrtc.SessionDescription{
					Type: webrtc.SDPTypeOffer,
					SDP:  string(sdp),
				}); err != nil {
					log.Printf("设置远程描述错误: %v", err)
					continue
				}
				log.Println("成功设置远程描述")
				remoteDescSet = true

				for _, candidate := range pendingCandidates {
					if err := peerConnection.AddICECandidate(candidate); err != nil {
						log.Printf("添加缓存的 ICE 候选者错误: %v", err)
					} else {
						log.Println("成功添加缓存的 ICE 候选者")
					}
				}
				pendingCandidates = nil

				log.Println("创建应答")
				answer, err := peerConnection.CreateAnswer(nil)
				if err != nil {
					log.Printf("创建应答错误: %v", err)
					continue
				}

				log.Println("设置本地描述")
				if err := peerConnection.SetLocalDescription(answer); err != nil {
					log.Printf("设置本地描述错误: %v", err)
					continue
				}

				answerSdp := answer.SDP
				logAnswer := answerSdp
				if len(answerSdp) > 200 {
					logAnswer = answerSdp[:200] + "..."
				}
				log.Printf("应答 SDP: %s", logAnswer)
				answerSDP := base64.StdEncoding.EncodeToString([]byte(answer.SDP))
				if err := conn.WriteJSON(Signal{
					SDP:  answerSDP,
					Type: "answer",
				}); err != nil {
					log.Printf("WebSocket 写入错误: %v", err)
					continue
				}
				log.Println("已发送应答给客户端")

				log.Println("等待 ICE 收集完成")
				gatherComplete := webrtc.GatheringCompletePromise(peerConnection)
				<-gatherComplete
				log.Println("ICE 收集完成")

				log.Println("启动视频流")
				go startVideo(ctx, videoTrack)
				log.Println("启动音频流")
				go startAudio(ctx, audioTrack)

			case "candidate":
				log.Println("收到客户端的 ICE 候选者")
				var candidate webrtc.ICECandidateInit
				if err := json.Unmarshal([]byte(signal.Candidate), &candidate); err != nil {
					log.Printf("ICE 候选者解析错误: %v", err)
					continue
				}
				if !remoteDescSet {
					log.Println("缓存 ICE 候选者，远程描述未设置")
					pendingCandidates = append(pendingCandidates, candidate)
				} else {
					if err := peerConnection.AddICECandidate(candidate); err != nil {
						log.Printf("添加 ICE 候选者错误: %v", err)
						continue
					}
					log.Println("成功添加 ICE 候选者")
				}
			default:
				log.Printf("未知消息类型: %s", signal.Type)
			}
		}
	}()

	wg.Wait()
}

func startVideo(ctx context.Context, videoTrack *webrtc.TrackLocalStaticSample) {
	log.Println("打开 H.264 文件进行视频流处理")
	file, err := os.Open("/home/weiqiangren/Documents/video.h264")
	if err != nil {
		log.Fatalf("打开 H.264 文件失败: %v", err)
	}
	defer file.Close()

	reader := bufio.NewReader(file)
	writer := rtpio.NewH264RTPWriter(videoTrack, 90000)
	log.Println("启动 H.264 视频流处理")
	startTime := time.Now()
	sampleCount := uint64(0)

	buf := make([]byte, 1024*1024) // 1MB 缓冲区
	for {
		select {
		case <-ctx.Done():
			log.Println("视频流被取消")
			return
		default:
			n, err := reader.Read(buf)
			if err == io.EOF {
				log.Println("视频流结束 (文件结束)")
				return
			}
			if err != nil {
				log.Printf("读取 H.264 文件错误: %v", err)
				return
			}

			data := buf[:n]
			_, err = writer.Write(data, uint64(sampleCount*40)) // 40ms per frame (25fps)
			if err != nil {
				log.Printf("写入 RTP 视频包错误: %v", err)
				return
			}
			log.Printf("写入 %d 字节到 WebRTC 视频轨道, 时间戳: %v", n, time.Duration(sampleCount)*40*time.Millisecond)

			sampleCount++
			elapsed := time.Since(startTime)
			expected := time.Duration(sampleCount) * 40 * time.Millisecond
			if elapsed < expected {
				time.Sleep(expected - elapsed)
			}
		}
	}
}

func startAudio(ctx context.Context, audioTrack *webrtc.TrackLocalStaticSample) {
	log.Println("打开 Opus 文件进行音频流处理")
	file, err := os.Open("/home/weiqiangren/Documents/audio.opus")
	if err != nil {
		log.Fatalf("打开 Opus 文件失败: %v", err)
	}
	defer file.Close()

	reader := bufio.NewReader(file)
	writer := rtpio.NewOpusRTPWriter(audioTrack, 48000)
	log.Println("启动 Opus 音频流处理")
	startTime := time.Now()
	sampleCount := uint64(0)

	buf := make([]byte, 1024*1024) // 1MB 缓冲区
	for {
		select {
		case <-ctx.Done():
			log.Println("音频流被取消")
			return
		default:
			n, err := reader.Read(buf)
			if err == io.EOF {
				log.Println("音频流结束 (文件结束)")
				return
			}
			if err != nil {
				log.Printf("读取 Opus 文件错误: %v", err)
				return
			}

			data := buf[:n]
			_, err = writer.Write(data, uint64(sampleCount*20)) // 20ms per frame
			if err != nil {
				log.Printf("写入 RTP 音频包错误: %v", err)
				return
			}
			log.Printf("写入 %d 字节到 WebRTC 音频轨道, 时间戳: %v", n, time.Duration(sampleCount)*20*time.Millisecond)

			sampleCount++
			elapsed := time.Since(startTime)
			expected := time.Duration(sampleCount) * 20 * time.Millisecond
			if elapsed < expected {
				time.Sleep(expected - elapsed)
			}
		}
	}
}
